/*******************************************************************************
 * Copyright (c) 2014 Greg Marut.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the GNU Public License v3.0
 * which accompanies this distribution, and is available at
 * http://www.gnu.org/licenses/gpl.html
 * 
 * Contributors:
 *    Greg Marut - initial API and implementation
 ******************************************************************************/
package com.gregmarut.network.tcp.secure;

import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.PBEParameterSpec;

import com.gregmarut.network.tcp.Client;
import com.gregmarut.util.ByteArray;
import com.gregmarut.util.Log;
import com.gregmarut.util.encryption.MD5;

/**
 * Provides the exact same functionality as a Client but also provides added
 * security by encrypting data that is transmitted over the sockets.
 * 
 * @author Greg Marut
 */
public abstract class SecureClient<S> extends Client<S>
{
	// ** Finals **//
	// holds the algorithm to use for this encryption
	public static final String ALGORITHM = "PBEWithMD5AndDES";
	public static final int ITERATION_COUNT = 1024;
	public static final int SALT_LENGTH = 8;
	
	// ** Objects **//
	// holds the cipher object responsible for encrypting and decrypting the data
	private final Cipher encryption;
	private final Cipher decryption;
	
	/**
	 * The constructor for SecureClient
	 * 
	 * @param host
	 * @param port
	 * @param encryptionKey
	 */
	public SecureClient(final String encryptionKey)
		throws InvalidEncryptionException
	{
		this(encryptionKey.getBytes());
	}
	
	/**
	 * The constructor for SecureClient
	 * 
	 * @param host
	 * @param port
	 * @param encryptionKey
	 */
	public SecureClient(final byte[] encryptionKey)
		throws InvalidEncryptionException
	{
		try
		{
			// create the secret key
			byte[] byteSalt = new ByteArray(MD5.encrypt(encryptionKey)).subArray(0, SALT_LENGTH).getBytes();
			PBEParameterSpec specParams = new PBEParameterSpec(byteSalt, ITERATION_COUNT);
			KeySpec keySpec = new PBEKeySpec(new String(encryptionKey).toCharArray());
			SecretKey key = SecretKeyFactory.getInstance(ALGORITHM).generateSecret(keySpec);
			
			// create the cipher object for encryption
			encryption = Cipher.getInstance(ALGORITHM);
			encryption.init(Cipher.ENCRYPT_MODE, key, specParams);
			
			// create the cipher object for decryption
			decryption = Cipher.getInstance(ALGORITHM);
			decryption.init(Cipher.DECRYPT_MODE, key, specParams);
		}
		catch (InvalidKeyException e)
		{
			// make sure the logger should display this error
			if (Log.DISPLAY_INTERNAL_ERRORS)
			{
				// log the error
				Log.err(e);
			}
			
			// throw a new invalid encryption exception
			throw new InvalidEncryptionException(e);
		}
		catch (InvalidKeySpecException e)
		{
			// make sure the logger should display this error
			if (Log.DISPLAY_INTERNAL_ERRORS)
			{
				// log the error
				Log.err(e);
			}
			
			// throw a new invalid encryption exception
			throw new InvalidEncryptionException(e);
		}
		catch (NoSuchAlgorithmException e)
		{
			// make sure the logger should display this error
			if (Log.DISPLAY_INTERNAL_ERRORS)
			{
				// log the error
				Log.err(e);
			}
			
			// throw a new invalid encryption exception
			throw new InvalidEncryptionException(e);
		}
		catch (NoSuchPaddingException e)
		{
			// make sure the logger should display this error
			if (Log.DISPLAY_INTERNAL_ERRORS)
			{
				// log the error
				Log.err(e);
			}
			
			// throw a new invalid encryption exception
			throw new InvalidEncryptionException(e);
		}
		catch (InvalidAlgorithmParameterException e)
		{
			// make sure the logger should display this error
			if (Log.DISPLAY_INTERNAL_ERRORS)
			{
				// log the error
				Log.err(e);
			}
			
			// throw a new invalid encryption exception
			throw new InvalidEncryptionException(e);
		}
	}
	
	/**
	 * Writes the data to the output stream
	 * 
	 * @param objData
	 * @throws IOException
	 */
	@Override
	protected void writeData(final byte[] decryptedData) throws IOException
	{
		byte[] encryptedData;
		
		try
		{
			//encrypt the data
			encryptedData = encryption.doFinal(decryptedData);
			
			// write the data to the output stream
			out.writeFrame(encryptedData);
		}
		catch (IllegalBlockSizeException e)
		{
			throw new IOException(e);
		}
		catch (BadPaddingException e)
		{
			throw new IOException(e);
		}
	}
	
	/**
	 * Returns the client activity to use for this client
	 * 
	 * @return
	 */
	@Override
	protected ClientActivity getClientActivity()
	{
		return new SecureClientActivity();
	}
	
	/**
	 * A class used to handle the connection process to the server
	 * 
	 * @author Greg
	 */
	protected class SecureClientActivity extends Client<S>.ClientActivity
	{
		/**
		 * Waits for data to be read from the reader
		 * 
		 * @return
		 * @throws IOException
		 */
		@Override
		protected byte[] readData() throws IOException
		{
			// holds the encrypted data
			byte[] encryptedData = super.readData();
			
			//make sure the data is not null
			if(null != encryptedData)
			{
				try
				{
					//decrypt the data
					byte[] decryptedData = decryption.doFinal(encryptedData);
					
					return decryptedData;
				}
				catch (IllegalBlockSizeException e)
				{
					throw new IOException(e);
				}
				catch (BadPaddingException e)
				{
					throw new IOException(e);
				}
			}
			else
			{
				return null;
			}
		}
	}
}
